import paramiko
import os
import sys
import stat  # 用于检查文件是否是目录

def create_sftp_client(host, port, username, password):
    """
    创建 SFTP 客户端以连接到远程服务器
    """
    client = paramiko.SSHClient()
    client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    client.connect(hostname=host, port=port, username=username, password=password)
    return client.open_sftp()

def progress(transferred, total):
    """
    显示文件下载进度
    """
    percentage = (transferred / total) * 100
    sys.stdout.write(f"\rDownloaded: {percentage:.2f}%")
    sys.stdout.flush()

def download_folder(sftp, remote_dir, local_dir):
    """
    递归下载远程文件夹到本地文件夹
    """
    if not os.path.exists(local_dir):
        os.makedirs(local_dir)

    for item in sftp.listdir_attr(remote_dir):
        remote_path = f"{remote_dir}/{item.filename}"
        local_path = os.path.join(local_dir, item.filename)

        # 使用 stat.S_ISDIR 检查是否是目录
        if stat.S_ISDIR(item.st_mode):  # 如果是文件夹，递归下载
            print(f"Entering directory: {remote_path}")
            download_folder(sftp, remote_path, local_path)
        else:  # 如果是文件，下载文件
            print(f"Downloading file: {remote_path} to {local_path}")
            sftp.get(remote_path, local_path, callback=lambda x, y: progress(x, y))
            print("\nDownload complete for file:", item.filename)

def upload_file(sftp, local_path, remote_path):
    """
    上传单个文件到远程路径
    """
    print(f"Uploading file: {local_path} to {remote_path}")
    sftp.put(local_path, remote_path, callback=lambda x, y: progress(x, y))
    print("\nUpload complete for file:", os.path.basename(local_path))

def upload_folder(sftp, local_dir, remote_dir):
    """
    递归上传本地文件夹到远程服务器
    """
    # 检查远程目录是否存在，如果不存在则创建
    try:
        sftp.stat(remote_dir)
    except FileNotFoundError:
        print(f"Creating remote directory: {remote_dir}")
        sftp.mkdir(remote_dir)

    for item in os.listdir(local_dir):
        local_path = os.path.join(local_dir, item)
        remote_path = f"{remote_dir}/{item}"

        if os.path.isdir(local_path):  # 如果是文件夹，递归上传
            print(f"Entering directory: {local_path}")
            upload_folder(sftp, local_path, remote_path)
        else:  # 如果是文件，上传文件
            upload_file(sftp, local_path, remote_path)


if __name__ == "__main__":
    # 配置远程服务器信息
    host = "connect.nma1.seetacloud.com"
    port = 18471  # 默认 SFTP 端口
    username = "root"
    password = "0f43xrZ5PtJa"

    # 远程和本地路径
    remote_dir = "/root/autodl-tmp/llama_vl_train/saved_model/robovlm"
    local_dir = "/Users/yangzhikai/Documents/code/robovlm/checkpoints"
    # 创建 SFTP 客户端并下载文件夹
    try:
        sftp_client = create_sftp_client(host, port, username, password)
        print(f"Starting download of folder: {remote_dir}")
        download_folder(sftp_client, remote_dir, local_dir)
        print("Upload complete!")
    except Exception as e:
        print(f"An error occurred: {e}")
    finally:
        sftp_client.close()
